package org.nlp.algo.classifier;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;

import com.aliasi.classify.Classification;
import com.aliasi.classify.Classified;
import com.aliasi.classify.TfIdfClassifierTrainer;
import com.aliasi.tokenizer.TokenFeatureExtractor;
import com.aliasi.util.Files;

/**
 * 使用 Lingpipe的TF/IDF分类器训练语料
 * 
 */
public class TrainTClassifier {
	
	// 语料文件夹
	private static String corpusPath = "resource/corpus";
		
	// 训练语料文件
	private static File trainDir = new File(corpusPath + "/train");
	
	//模型文件名
	private static String modelFile = "resource/model/tclassifier.model";
	
	//定义分类
	private static String[] CATEGORIES;// = { "汽车", "财经", "IT", "健康", "体育", "旅游", "教育", "招聘", "文化", "军事"};

	public static void main(String[] args) throws ClassNotFoundException,
			IOException {
		
//		TokenizerFactory tokenizerFactory = LuceneTokenizerFactory.INSTANCE;
		TfIdfClassifierTrainer<CharSequence> classifier = new TfIdfClassifierTrainer<CharSequence>(
				new TokenFeatureExtractor(LuceneTokenizerFactory.INSTANCE));		
//		TfIdfClassifierTrainer<CharSequence> classifier = new TfIdfClassifierTrainer<CharSequence>(
//				new TokenFeatureExtractor(CharacterTokenizerFactory.INSTANCE));

		// 开始训练
		for (int i = 0; i < CATEGORIES.length; i++) {
			File classDir = new File(trainDir, CATEGORIES[i]);
			if (!classDir.isDirectory()) {
				System.out.println("不能找到目录=" + classDir);
			}

			// 训练器遍历分类文件夹下的所有文件
			for (File file : classDir.listFiles()) {
				String text = Files.readFromFile(file, "gbk");	//训练文件gbk编码							
				System.out.println("正在训练 " + CATEGORIES[i] + file.getName());
				Classification classification = new Classification(
						CATEGORIES[i]);
				Classified<CharSequence> classified = new Classified<CharSequence>(
						text, classification);
				classifier.handle(classified);
			} 
		}
		

		// 把分类器模型写到文件上
		System.out.println("开始生成分类器");
		
		ObjectOutputStream os = new ObjectOutputStream(new FileOutputStream(
				modelFile));
		classifier.compileTo(os);
		os.close();
		
		System.out.println("分类器生成完成");
	}
}